import sys
import os
from typing import List, Dict, Tuple, Union
import json

DEBUG = False
# 初始化指令集映射
opcodes = {}
# 操作码映射
calc_ops = {}
logic_ops = {}
set_ops = {}
fop_ops = {}
labels = {}


def transImm(imm: str) -> int:
    if imm.startswith('0x'):
        return int(imm, 16)
    elif imm.startswith('0b'):
        return int(imm, 2)
    elif imm.startswith('0o'):
        return int(imm, 8)
    elif imm.startswith('@'):
        return labels[imm[1:]]
    else:
        return int(imm)


class Data:
    def __init__(self, data: str):
        self.src: str = data
        self.type: str = ""
        self.addr: int = 0
        self.values: List[int] | str = []

    def __call__(self):
        parts = self.src.split(maxsplit=3)
        if len(parts) < 4:
            raise ValueError(f"Invalid data definition: {self.src}")

        self.type = parts[0]
        self.addr = int(parts[2], 16)
        self.values = parts[-1]

        match self.type:
            case '.byte' | '.2byte' | '.4byte' | '.8byte':
                self.values = [int(i, 16) if i.startswith(
                    '0x') else int(i) for i in self.values.split(',')]
            case '.string':
                # 移除字符串的引号
                if (self.values.startswith('"') and self.values.endswith('"')) or \
                   (self.values.startswith("'") and self.values.endswith("'")):
                    self.values = self.values[1:-1]
                    self.values = self.values.replace('\\n', '\n')
        return self

    def to_bytes(self):
        result = {'addr': self.addr, 'data': b''}
        match self.type:
            case '.byte':
                for val in self.values:
                    result['data'] += val.to_bytes(1, 'little')
            case '.2byte':
                for val in self.values:
                    result['data'] += val.to_bytes(2, 'little')
            case '.4byte':
                for val in self.values:
                    result['data'] += val.to_bytes(4, 'little')
            case '.8byte':
                for val in self.values:
                    result['data'] += val.to_bytes(8, 'little')
            case '.string':
                result['data'] += self.values.encode('utf-8') + b'\0'
        return result


class Command:
    def __init__(self, command: str, count: int):
        self.src: str = command
        self.count: int = count
        self.opcode: str | int = 0
        self.args: List[str] = []
        self.r1: int = 0
        self.r2: int = 0
        self.r3: int = 0
        self.imm: int = 0

    def to_bytes(self):
        if isinstance(self.opcode, str):
            self.opcode = int(opcodes[self.opcode], 16)
        instruction = (self.imm << 22) | (self.r3 << 17) | (
            self.r2 << 12) | (self.r1 << 7) | self.opcode
        instruction &= 0xffff_ffff
        return instruction.to_bytes(4, 'little')

    def __call__(self):
        # Split the command into opcode and arguments
        parts = self.src.split(maxsplit=1)
        self.opcode = parts[0]
        # Split arguments by commas and strip whitespace
        self.args = [arg.strip() for arg in parts[1].split(',')
                     ] if len(parts) > 1 else []
        match self.opcode:
            case 'nop':
                self.r1, self.r2, self.r3, self.imm = 0, 0, 0, 0
            case 'iret' | 'halt':
                self.r1, self.r2, self.r3, self.imm = 31, 31, 31, -1
            case 'trap':
                self.imm = int(self.args[0], 16)
                self.r1 = self.r2 = self.r3 = 0
            case 'calc' | 'logic' | 'fop' | 'dop':
                self.r1 = int(self.args[0][1:])
                self.r2 = int(self.args[1][1:])
                self.r3 = int(self.args[2][1:])
                if self.opcode == 'calc':
                    self.imm = calc_ops[self.args[3]]
                elif self.opcode == 'logic':
                    self.imm = logic_ops[self.args[3]]
                else:
                    self.imm = fop_ops[self.args[3]]
            case 'set':
                self.r1 = int(self.args[0][1:])
                try:  # set r1,0x201
                    self.imm = transImm(self.args[1])
                    self.r2 = set_ops["mov"]
                except ValueError:  # set r1,r2<op>0x201
                    arg = self.args[-1]
                    if arg.find('+') != -1:  # set r1,r2+0x201
                        self.r2 = set_ops['add']
                        arg = arg.split('+')
                        self.r3 = int(arg[0][1:])
                        self.imm = transImm(arg[1])
                    elif arg.find('-') != -1:
                        self.r2 = set_ops['sub']
                        arg = arg.split('-')
                        self.r3 = int(arg[0][1:])
                        self.imm = transImm(arg[1])
                    elif arg.find('<<') != -1:
                        self.r2 = set_ops['lsh']
                        arg = arg.split('<<')
                        self.r3 = int(arg[0][1:])
                        self.imm = int(arg[1])
                    elif arg.find('>>') != -1:
                        self.r2 = set_ops['rsh']
                        arg = arg.split('>>')
                        self.r3 = int(arg[0][1:])
                        self.imm = int(arg[1])
                    else:
                        raise ValueError(f"Invalid set operation: {arg}")
            case 'load' | 'store' | 'loadf' | 'storef':
                arg = self.args[2].split('+')
                self.r1 = int(self.args[0][1:])
                self.r2 = int(arg[0][1:])
                self.r3 = int(self.args[1])
                if arg[1].startswith('0x'):
                    self.imm = int(arg[1], 16)
                else:
                    self.imm = int(arg[1])
            case 'jmp' | 'call':
                if len(self.args) == 1:  # condition-less jump/call
                    self.r1 = 0
                    self.r3 = labels['all']
                else:  # conditioned jump/call
                    self.r1 = int(self.args[0][1:])
                    self.r3 = labels[self.args[1]]
                if self.args[-1].find('+') != -1:  # long jump/call
                    self.r3 |= labels['abs']
                    arg = self.args[-1].split('+')
                    self.r2 = int(arg[0][1:])
                    self.imm = transImm(arg[1])
                else:  # short jump/call
                    self.r2 = 0
                    self.imm = transImm(self.args[-1])
            case 'ret':
                self.r1, self.r2, self.r3, self.imm = 0, labels['all'], 0, 0
            case 'kload' | 'kstore' | 'read' | 'write':
                self.r1 = int(self.args[0][1:])
                self.r2 = int(self.args[1])
                if len(self.args) > 2:
                    self.imm = transImm(self.args[2])
        return self


class Sector:
    def __init__(self):
        self.name = ""
        self.beg_addr = 0
        self.size = 0
        self.commands: List[Command | Data] = []
        self.data: List[bytes] = []

    def __str__(self):
        ret = f"name: {self.name}\t"
        ret += f"beg_addr: {self.beg_addr}\t"
        ret += f"size: {self.size}\t"
        ret += f"commands: \n"
        for i in self.commands:
            ret += f"{i.src}\n"
        return ret

    def __call__(self):
        self.commands = [i() for i in self.commands]
        self.data = []
        for cmd in self.commands:
            try:
                self.data.append(cmd.to_bytes())
            except ValueError as e:
                print(f"Error in command '{cmd.src}': {e}")
                raise e
            except Exception as e:
                print(f"Unexpected error in command '{cmd.src}': {e}")
                raise e
        ret = b""
        for i in self.data:
            if isinstance(i, bytes):
                ret += i
            elif isinstance(i, dict):
                if len(ret) < (i['addr']-self.beg_addr):
                    ret += b'\0' * ((i['addr']-self.beg_addr) - len(ret))
                ret += i['data']
        return ret


class Assembler:
    def __init__(self, input_file: str, output_file: str):
        self.src = input_file
        self.dst = output_file
        self.sectors: List[Sector] = []
        self.data = []

    def __call__(self):
        with open(self.src, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        lines = [i.split('#')[0].strip('\n').strip() for i in lines]
        lines = [i for i in lines if i != '']
        sections = {}  # 存储段名称和起始地址的映射
        current_section = None
        count = 0
        # 第一遍处理：收集段信息和指令计数
        for line in lines:
            line = line.strip()
            if not line:
                continue
            if line.startswith('.'):
                # 处理段定义行（如：.section .text at 0）
                if line.startswith('.section'):
                    parts = line.split()
                    section_name = parts[1]
                    section_addr = int(parts[3][:-1], 0)  # 支持16进制/10进制地址
                    sections[section_name] = section_addr
                    labels[section_name] = section_addr
                    current_section = section_name
                # 更新段偏移量
                count = 0
                continue
            if line.endswith(':'):
                # 跳过标记定义行
                continue
            # 统计实际指令
            count += 1
        # 第二遍处理：记录标签地址（包含段偏移）
        current_section = None
        count = 0
        for line in lines:
            line = line.strip()
            if not line:
                continue
            if line.startswith('.'):
                # 更新当前段
                if line.startswith('.section'):
                    current_section = line.split()[1]
                count = 0
                continue
            if line.endswith(':'):
                # 记录标签地址（段起始地址 + 段内偏移）
                label = line[:-1]
                labels[label] = sections.get(current_section, 0) + count
                continue
            # 统计实际指令
            count += 1
        # Second pass: process commands
        sector = None
        count = 0
        for i in lines:
            if i.startswith('.'):
                vals = i.split(' ')
                vals = [j for j in vals if (j != ' ' and j != '')]
                match vals[0]:
                    case '.section':
                        if sector is not None:
                            self.sectors.append(sector)
                        sector = Sector()
                        sector.name = vals[1]
                        count = sections[sector.name]
                        if len(vals) == 4:
                            sector.beg_addr = int(vals[-1][:-1], 16)
                        else:
                            sector.beg_addr = 0
                    case '.global':
                        if 'at' in vals:
                            labels[vals[1]] = int(vals[-1], 16)
                        else:
                            pass
                    case '.equ':
                        try:
                            labels[vals[1]] = int(vals[2])
                        except ValueError:
                            labels[vals[1]] = int(vals[2], 16)
                    case '.string' | '.byte' | '.2byte' | '.4byte' | '.8byte':
                        labels[vals[1]] = int(vals[2], 16)
                        if sector:
                            sector.commands.append(Data(i))
            else:
                if sector is None:
                    sector = Sector()
                    sector.name = '.text'
                if i[-1] != ':':
                    count += 1
                    sector.commands.append(Command(i, count))
        if sector is not None:
            self.sectors.append(sector)
        self.data = [i() for i in self.sectors]
        with open(f"output/{self.dst}", 'wb') as f:
            for i in range(len(self.data)):
                while f.seek(0, os.SEEK_END) < self.sectors[i].beg_addr:
                    f.write(b'\0'*(512-f.tell()))
                f.seek(self.sectors[i].beg_addr)
                f.write(self.data[i])


if __name__ == "__main__":
    dir_path = sys.argv[3]
    with open(os.path.join(dir_path, 'command.json'), 'r') as f:
        opcodes = json.load(f)
    with open(os.path.join(dir_path, 'ops.json'), 'r') as f:
        op_list = json.load(f)
    calc_ops = op_list['calc_ops']
    logic_ops = op_list['logic_ops']
    fop_ops = op_list['fop_ops']
    set_ops = op_list['set_ops']
    labels = op_list['labels']
    assembler = Assembler(sys.argv[1], sys.argv[2])
    # assembler = Assembler('Processor/assemble/test_io.vasm', 'test_io.bin')
    assembler()
